from ultralytics import YOLO

if __name__ == '__main__':
    # 加载模型
    model = YOLO('ultralytics/cfg/models/v8/yolov8.yaml')
    # 训练参数(根据论文设置)
    model.train(
        data='data.yaml',
        epochs=200,
        batch=32,  # 降低batch size避免内存不足
        imgsz=512,  # 使用标准输入尺寸
        device=0,  # GPU
        workers=0,  # 单进程模式避免Windows多进程问题

        # 优化器设置
        optimizer='SGD',
        lr0=0.01,
        momentum=0.937,
        weight_decay=0.0005,

        # 数据增强(论文中提到的)
        mosaic=1.0,
        mixup=0.0,
        hsv_h=0.015,  # 色调
        hsv_s=0.7,  # 饱和度
        hsv_v=0.4,  # 明度
        degrees=0.0,
        translate=0.1,  # 随机平移
        scale=0.5,
        flipud=0.0,
        fliplr=0.5,

        # 其他设置
        patience=50,
        save=True,
        save_period=50,
        plots=True,
        verbose=True,

        # 项目名称
        project='runs/train',
        name='leeyolo-crack',
    )

    print("训练完成!")